package cmd

import (
	"context"
	"fmt"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer/roundrobin"
	"google.golang.org/grpc/resolver"
	"google.golang.org/grpc/resolver/manual"

	"gitee.com/zfd81/dbrs/engine"
	"gitee.com/zfd81/dbrs/internal/token"
	pb "gitee.com/zfd81/dbrs/proto/dbrspb"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

var ()

func NewTestCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "test",
		Short: "Test databroker connection",
		Run:   testCommandFunc,
	}

	//设置命令行参数
	cmd.Flags().StringVar(&address, "addr", conf.Address, "Network address to run the server")
	cmd.Flags().Uint32Var(&port, "p", conf.Port, "Port to run the server")

	return cmd
}

func testCommandFunc(cmd *cobra.Command, args []string) {
	claims := token.BrokerClaims{
		Tenant: "testTenant",
		Sign:   "testSign",
		User:   "testUser",
	}
	tokenString, err := token.Create(claims, engine.Token_Secret, engine.Token_Minute)
	if err != nil {
		logrus.Errorf("Create token error: %v\n", err)
		return
	}

	for _, database := range conf.Databases {
		client, err := GetClient(database.Endpoints)
		if err != nil {
			logrus.Errorf("Test database %s connection error: %v\n", database.Name, err)
		} else {
			_, err = client.Test(context.Background(), &pb.TestRequest{
				Token: tokenString,
			})
			if err != nil {
				logrus.Errorf("Test error: %v\n", err)
			} else {
				logrus.Infof("Test database %s connection succeeded.\n", database.Name)
			}
		}

	}
}

func GetClient(endpoints []string) (pb.BrokerClient, error) {
	var addresses []resolver.Address
	for _, endpoint := range endpoints {
		addresses = append(addresses, resolver.Address{Addr: endpoint})
	}

	r := manual.NewBuilderWithScheme("whatever")
	r.InitialState(resolver.State{Addresses: addresses})
	conn, err := grpc.Dial(
		r.Scheme()+":///grpcs.server",
		grpc.WithInsecure(),
		grpc.WithResolvers(r),
		grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"LoadBalancingPolicy": "%s", "RetryPolicy": {"MaxAttempts":2, "InitialBackoff": "0.1s", "MaxBackoff": "1s", "BackoffMultiplier": 2.0, "RetryableStatusCodes": ["UNAVAILABLE"]}}`, roundrobin.Name)),
		grpc.WithBlock(),
		grpc.WithTimeout(3*time.Second))
	if err != nil {
		return nil, fmt.Errorf("did not connect broker: %v", err)
	}
	return pb.NewBrokerClient(conn), nil
}
